import uuid
import logging
from typing import Optional
from fastapi import APIRouter, UploadFile, File, Form, HTTPException, BackgroundTasks, Request
from fastapi.responses import JSONResponse
from PIL import Image
import io
import os

from app.models.request import ColorizeRequest
from app.models.response import ColorizeResponse, TaskStatusResponse
from app.services.colorizer import colorization_service
from app.services.storage import storage_service
from app.core.config import settings

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/colorize", tags=["colorization"])

# In-memory task storage (replace with Redis in production)
task_storage = {}
MAX_BATCH_SIZE = 10
MAX_CONCURRENT_TASKS = 100


def validate_file(file: UploadFile) -> None:
    """Validate uploaded file"""
    if not file.content_type:
        raise HTTPException(
            status_code=400,
            detail="Content-Type header is missing"
        )
    
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file type: {file.content_type}. Only images are allowed."
        )
    
    # Check file extension
    allowed_extensions = settings.allowed_extensions
    file_extension = file.filename.split('.')[-1].lower()
    if file_extension not in allowed_extensions:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file extension: {file_extension}. Allowed: {allowed_extensions}"
        )


async def validate_image_content(content: bytes) -> Image.Image:
    """Validate image content and return PIL Image"""
    try:
        # Check file size
        if len(content) > settings.max_file_size:
            raise HTTPException(
                status_code=413,
                detail=f"File too large. Maximum size: {settings.max_file_size / 1024 / 1024:.1f}MB"
            )
        
        # Validate image format
        image = Image.open(io.BytesIO(content))
        
        # Verify image can be processed
        image.verify()
        
        # Reopen after verify
        image = Image.open(io.BytesIO(content))
        
        # Ensure RGB mode
        if image.mode not in ["RGB", "RGBA"]:
            image = image.convert("RGB")
        
        return image
        
    except Image.UnidentifiedImageError:
        raise HTTPException(
            status_code=400,
            detail="Invalid image format or corrupted image"
        )
    except Exception as e:
        raise HTTPException(
            status_code=400,
            detail=f"Image validation failed: {str(e)}"
        )


@router.post("/", response_model=ColorizeResponse)
async def colorize_image(
    background_tasks: BackgroundTasks,
    file: UploadFile = File(...),
    style: str = Form("anime"),
    strength: float = Form(0.7),
    guidance_scale: float = Form(7.5),
    num_inference_steps: int = Form(20),
    seed: Optional[int] = Form(None)
):
    """
    Upload a manga sketch and get it colorized
    
    - **file**: The manga sketch image file (PNG, JPG, JPEG, or WEBP)
    - **style**: Colorization style (anime, realistic, cartoon, sketch, watercolor)
    - **strength**: How much to modify the original (0.1-1.0)
    - **guidance_scale**: How closely to follow the prompt (1.0-20.0)
    - **num_inference_steps**: Number of denoising steps (10-50)
    - **seed**: Random seed for reproducible results
    """
    try:
        # Validate file
        validate_file(file)
        
        # Read file content
        content = await file.read()
        
        # Validate image content
        image = await validate_image_content(content)
        
        # Check concurrent task limit
        if len(task_storage) >= MAX_CONCURRENT_TASKS:
            raise HTTPException(
                status_code=429,
                detail="Too many concurrent tasks. Please try again later."
            )
        
        # Validate parameters
        if style not in ["anime", "realistic", "cartoon", "sketch", "watercolor"]:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid style: {style}"
            )
        
        strength = max(0.1, min(1.0, strength))
        guidance_scale = max(1.0, min(20.0, guidance_scale))
        num_inference_steps = max(10, min(50, num_inference_steps))
        
        # Generate task ID
        task_id = str(uuid.uuid4())
        
        # Save original image
        original_filename = f"original_{task_id}.png"
        original_path = await storage_service.save_file(
            content,
            original_filename,
            prefix="original"
        )
        
        # Create task
        task = {
            "task_id": task_id,
            "original_path": original_path,
            "original_filename": file.filename,
            "style": style,
            "strength": strength,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps,
            "seed": seed,
            "status": "pending",
            "created_at": time.time()
        }
        
        # Store task
        task_storage[task_id] = task
        
        # Queue the colorization task
        background_tasks.add_task(
            _process_colorization,
            task_id,
            content,
            style,
            strength,
            guidance_scale,
            num_inference_steps,
            seed
        )
        
        return ColorizeResponse(
            success=True,
            task_id=task_id,
            status="pending",
            original_image_url=storage_service.get_file_url(original_filename)
        )
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Colorization request failed: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Internal server error: {str(e)}"
        )


@router.get("/status/{task_id}", response_model=TaskStatusResponse)
async def get_task_status(task_id: str):
    """Get the status of a colorization task"""
    if task_id not in task_storage:
        raise HTTPException(
            status_code=404,
            detail="Task not found"
        )
    
    task = task_storage[task_id]
    
    # Clean up old completed tasks (older than 1 hour)
    current_time = time.time()
    if task["status"] in ["completed", "failed"] and current_time - task.get("created_at", 0) > 3600:
        task_storage.pop(task_id, None)
        raise HTTPException(
            status_code=404,
            detail="Task expired"
        )
    
    return TaskStatusResponse(
        task_id=task_id,
        status=task["status"],
        progress=task.get("progress", 0),
        result=task.get("result"),
        estimated_time_remaining=task.get("estimated_time_remaining")
    )


@router.get("/result/{task_id}", response_model=ColorizeResponse)
async def get_result(task_id: str):
    """Get the final result of a colorization task"""
    if task_id not in task_storage:
        raise HTTPException(
            status_code=404,
            detail="Task not found"
        )
    
    task = task_storage[task_id]
    
    if task["status"] != "completed":
        raise HTTPException(
            status_code=202,
            detail="Task not completed yet"
        )
    
    return task["result"]


@router.post("/batch")
async def batch_colorize(
    background_tasks: BackgroundTasks,
    files: list[UploadFile] = File(...),
    style: str = Form("anime"),
    strength: float = Form(0.7),
    guidance_scale: float = Form(7.5),
    num_inference_steps: int = Form(20)
):
    """Colorize multiple images in batch"""
    if len(files) > MAX_BATCH_SIZE:
        raise HTTPException(
            status_code=413,
            detail=f"Maximum {MAX_BATCH_SIZE} files allowed in batch"
        )
    
    if len(files) == 0:
        raise HTTPException(
            status_code=400,
            detail="No files provided"
        )
    
    batch_id = str(uuid.uuid4())
    task_ids = []
    
    for file in files:
        try:
            validate_file(file)
            content = await file.read()
            image = await validate_image_content(content)
            
            task_id = str(uuid.uuid4())
            original_filename = f"original_{task_id}.png"
            original_path = await storage_service.save_file(
                content,
                original_filename,
                prefix="original"
            )
            
            task = {
                "task_id": task_id,
                "original_path": original_path,
                "original_filename": file.filename,
                "style": style,
                "strength": strength,
                "guidance_scale": guidance_scale,
                "num_inference_steps": num_inference_steps,
                "seed": None,
                "status": "pending",
                "created_at": time.time(),
                "batch_id": batch_id
            }
            
            task_storage[task_id] = task
            task_ids.append(task_id)
            
            background_tasks.add_task(
                _process_colorization,
                task_id,
                content,
                style,
                strength,
                guidance_scale,
                num_inference_steps,
                None
            )
            
        except Exception as e:
            logger.error(f"Failed to process file {file.filename}: {str(e)}")
            continue
    
    return {
        "batch_id": batch_id,
        "task_ids": task_ids,
        "message": f"Batch processing started with {len(task_ids)} tasks"
    }


@router.delete("/task/{task_id}")
async def delete_task(task_id: str):
    """Delete a task and associated files"""
    if task_id not in task_storage:
        raise HTTPException(
            status_code=404,
            detail="Task not found"
        )
    
    task = task_storage[task_id]
    
    try:
        # Delete original file
        if "original_path" in task:
            await storage_service.delete_file(task["original_path"])
        
        # Delete colored file if exists
        if task.get("result") and task["result"].get("colored_image_url"):
            # Extract filename from URL
            colored_filename = task["result"]["colored_image_url"].split("/")[-1]
            colored_path = f"colored_{colored_filename}"
            await storage_service.delete_file(colored_path)
        
        # Remove from storage
        task_storage.pop(task_id, None)
        
        return {"success": True, "message": "Task deleted successfully"}
        
    except Exception as e:
        logger.error(f"Failed to delete task {task_id}: {str(e)}")
        raise HTTPException(
            status_code=500,
            detail=f"Failed to delete task: {str(e)}"
        )


async def _process_colorization(
    task_id: str,
    image_bytes: bytes,
    style: str,
    strength: float,
    guidance_scale: float,
    num_inference_steps: int,
    seed: Optional[int]
):
    """Process the colorization task with comprehensive error handling"""
    try:
        if task_id not in task_storage:
            logger.error(f"Task {task_id} not found")
            return
        
        task = task_storage[task_id]
        task["status"] = "processing"
        task["progress"] = 10
        
        # Check if model is loaded
        if not colorization_service.is_model_loaded():
            raise RuntimeError("AI model not available")
        
        # Colorize the image
        task["progress"] = 30
        colored_bytes, metadata = colorization_service.colorize_sync(
            image_bytes,
            style=style,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            seed=seed
        )
        
        task["progress"] = 70
        
        # Save colored image
        colored_filename = f"colored_{task_id}.png"
        colored_path = await storage_service.save_file(
            colored_bytes,
            colored_filename,
            prefix="colored"
        )
        
        task["progress"] = 90
        
        # Create response
        response = ColorizeResponse(
            success=True,
            task_id=task_id,
            status="completed",
            colored_image_url=storage_service.get_file_url(colored_filename),
            processing_time=metadata["processing_time"],
            metadata=metadata
        )
        
        task["result"] = response
        task["status"] = "completed"
        task["progress"] = 100
        
        logger.info(f"Task {task_id} completed successfully")
        
    except Exception as e:
        logger.error(f"Task {task_id} failed: {str(e)}")
        if task_id in task_storage:
            task = task_storage[task_id]
            task["status"] = "failed"
            task["error_message"] = str(e)
            task["progress"] = 100


@router.get("/tasks")
async def list_tasks(
    status: Optional[str] = None,
    limit: int = 50,
    offset: int = 0
):
    """List tasks with filtering and pagination"""
    if limit > 100:
        raise HTTPException(
            status_code=400,
            detail="Limit cannot exceed 100"
        )
    
    tasks = []
    for task_id, task in list(task_storage.items()):
        if status is None or task["status"] == status:
            tasks.append({
                "task_id": task_id,
                "status": task["status"],
                "original_filename": task.get("original_filename"),
                "created_at": task.get("created_at"),
                "progress": task.get("progress", 0)
            })
    
    # Sort by creation time (newest first)
    tasks.sort(key=lambda x: x.get("created_at", 0), reverse=True)
    
    # Apply pagination
    paginated_tasks = tasks[offset:offset + limit]
    
    return {
        "tasks": paginated_tasks,
        "total": len(tasks),
        "offset": offset,
        "limit": limit
    }
